# 实验2~实验3的train和test数据划分

import os
from random import shuffle
import random
from util import *  

path_pad = 'E:/data/palm_vein/VERA-PAD'
path_real = os.path.join(path_pad, 'Real')
path_fake = os.path.join(path_pad, 'Spoof')
            
train_csv_contents = []
people_ids = list(range(1,50+1))
random.seed(123)  
shuffle(people_ids)
common_train_ids = people_ids[:25] # 编号1~50随机选一半
common_test_ids = people_ids[25:] # 编号1~50随机选另一半

# 真手掌类编号1~50
real_file_lst = recursive_files(path_real, path_pad)
cid = 1
pid_cid_dict = {}
real_file_lst_train = []
real_file_lst_test = []
for file in real_file_lst:
    pid = get_people_id(file)
    if pid in common_train_ids:
        real_file_lst_train.append(file)
        if pid not in pid_cid_dict:
            pid_cid_dict[pid] = cid
            cid += 1
            
        if 'L' in file:
            train_csv_contents.append([file, pid_cid_dict[pid]])
        else:
            train_csv_contents.append([file, pid_cid_dict[pid] + 25])
    else:
        real_file_lst_test.append(file)


# 假手掌类编号51~100
fake_file_lst = recursive_files(path_fake, path_pad)
fake_file_lst_train = []
fake_file_lst_test = []
for file in fake_file_lst:
    pid = get_people_id(file)
    if pid in common_train_ids:
        fake_file_lst_train.append(file)
        if 'L' in file:
            train_csv_contents.append([file, pid_cid_dict[pid] + 50])
        else:
            train_csv_contents.append([file, pid_cid_dict[pid] + 75])
    else:
        fake_file_lst_test.append(file)

# 保存train csv
write_csv('exp2_vera_pad_train.csv', train_csv_contents)



    
# 保存test 1 csv：全部用真手掌图像
pair_set = set()
# 类内500， 用1表示
i = 1
test_len = len(real_file_lst_test)
normal_test_csv_contents = []
while i<=500:
    pair_1 = random.randint(0, test_len-1)
    pair_2 = pair_1 + random.randint(-5, 5)
    if pair_1 == pair_2:
        continue

    if pair_2 < 0:
        pair_2 = 0
    if pair_2 >= test_len:
        pair_2 = test_len - 1
    
    if intra_class(real_file_lst_test[pair_1], real_file_lst_test[pair_2]):
        if (pair_1, pair_2) in pair_set or (pair_2, pair_1) in pair_set:
            pass
        else:
            pair_set.add((pair_1, pair_2))
            pair_set.add((pair_2, pair_1))
            normal_test_csv_contents.append([real_file_lst_test[pair_1], real_file_lst_test[pair_2], 1])
            i += 1

# 类间500，用0表示
i = 1
pair_set = set()
while i<=500:
    pair_1 = random.randint(0, test_len-1)
    pair_2 = random.randint(0, test_len-1)
    if pair_1 == pair_2:
        continue
    
    if not intra_class(real_file_lst_test[pair_1], real_file_lst_test[pair_2]):
        if (pair_1, pair_2) in pair_set or (pair_2, pair_1) in pair_set:
            pass
        else:
            pair_set.add((pair_1, pair_2))
            pair_set.add((pair_2, pair_1))
            normal_test_csv_contents.append([real_file_lst_test[pair_1], real_file_lst_test[pair_2], 0])
            i += 1

write_csv('exp2_vera_pad_test1.csv', normal_test_csv_contents)